package com.mini.grpc.common;

import com.mini.grpc.serialization.Marshaller;
import java.io.ByteArrayInputStream;
import java.io.InputStream;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;

/**
 * Metadata用于在RPC调用中传递元数据信息。
 * 
 * <p>元数据是一组键值对，键是ASCII字符串，值可以是字符串或二进制数据。
 */
public final class Metadata {
    
    /**
     * ASCII字符串编组器
     */
    public static final Marshaller<String> ASCII_STRING_MARSHALLER = new Marshaller<String>() {
        @Override
        public InputStream stream(String value) {
            return new ByteArrayInputStream(value.getBytes(StandardCharsets.UTF_8));
        }
        
        @Override
        public String parse(InputStream stream) {
            try {
                return new String(stream.readAllBytes(), StandardCharsets.UTF_8);
            } catch (Exception e) {
                throw new RuntimeException("Failed to parse ASCII string", e);
            }
        }
    };
    
    /**
     * 二进制字节编组器
     */
    public static final Marshaller<byte[]> BINARY_BYTE_MARSHALLER = new Marshaller<byte[]>() {
        @Override
        public InputStream stream(byte[] value) {
            return new ByteArrayInputStream(value);
        }
        
        @Override
        public byte[] parse(InputStream stream) {
            try {
                return stream.readAllBytes();
            } catch (Exception e) {
                throw new RuntimeException("Failed to parse binary bytes", e);
            }
        }
    };
    
    /**
     * 元数据键的后缀，表示键对应的值是二进制数据
     */
    public static final String BINARY_HEADER_SUFFIX = "-bin";
    
    /**
     * 元数据键的后缀，表示键对应的值是ASCII字符串
     */
    public static final String ASCII_HEADER_SUFFIX = "";
    
    /**
     * 元数据键的类型
     *
     * @param <T> 值的类型
     */
    public static class Key<T> {
        private final String name;
        private final boolean isBinary;
        private final Marshaller<T> marshaller;
        
        private Key(String name, boolean isBinary, Marshaller<T> marshaller) {
            this.name = name;
            this.isBinary = isBinary;
            this.marshaller = marshaller;
        }
        
        /**
         * 获取键名
         *
         * @return 键名
         */
        public String name() {
            return name;
        }
        
        /**
         * 判断键是否对应二进制数据
         *
         * @return 如果键对应二进制数据则返回true，否则返回false
         */
        public boolean isBinary() {
            return isBinary;
        }
        
        /**
         * 获取编组器
         *
         * @return 编组器
         */
        public Marshaller<T> getMarshaller() {
            return marshaller;
        }
        
        @Override
        public boolean equals(Object obj) {
            if (this == obj) {
                return true;
            }
            if (obj == null || getClass() != obj.getClass()) {
                return false;
            }
            Key<?> other = (Key<?>) obj;
            return name.equals(other.name) && isBinary == other.isBinary;
        }
        
        @Override
        public int hashCode() {
            return name.hashCode() ^ (isBinary ? 1 : 0);
        }
        
        @Override
        public String toString() {
            return "Key{name=" + name + ", isBinary=" + isBinary + "}";
        }
        
        /**
         * 创建一个键
         *
         * @param name 键名
         * @param marshaller 编组器
         * @return 新的Key实例
         */
        @SuppressWarnings("unchecked")
        public static <T> Key<T> of(String name, Marshaller<T> marshaller) {
            boolean isBinary = name.endsWith(BINARY_HEADER_SUFFIX);
            return new Key<>(name, isBinary, marshaller);
        }
        
        /**
         * 创建一个ASCII字符串类型的键
         *
         * @param name 键名
         * @return 新的Key实例
         */
        @SuppressWarnings("unchecked")
        public static Key<String> of(String name) {
            return new Key<>(name, false, (Marshaller<String>) ASCII_STRING_MARSHALLER);
        }
        
        /**
         * 创建一个二进制数据类型的键
         *
         * @param name 键名
         * @return 新的Key实例
         */
        @SuppressWarnings("unchecked")
        public static Key<byte[]> ofBinary(String name) {
            return new Key<>(name, true, (Marshaller<byte[]>) BINARY_BYTE_MARSHALLER);
        }
    }
    
    private final Map<String, List<Object>> store = new HashMap<>();
    
    /**
     * 创建一个空的元数据实例
     */
    public Metadata() {
    }
    
    /**
     * 获取键对应的所有值
     *
     * @param key 键
     * @param <T> 值的类型
     * @return 键对应的所有值，如果键不存在则返回空列表
     */
    @SuppressWarnings("unchecked")
    public <T> List<T> getAll(Key<T> key) {
        List<Object> values = store.get(key.name());
        if (values == null) {
            return Collections.emptyList();
        }
        List<T> result = new ArrayList<>(values.size());
        for (Object value : values) {
            result.add((T) value);
        }
        return result;
    }
    
    /**
     * 获取键对应的第一个值
     *
     * @param key 键
     * @param <T> 值的类型
     * @return 键对应的第一个值，如果键不存在则返回null
     */
    @SuppressWarnings("unchecked")
    public <T> T get(Key<T> key) {
        List<Object> values = store.get(key.name());
        if (values == null || values.isEmpty()) {
            return null;
        }
        return (T) values.get(0);
    }
    
    /**
     * 添加一个键值对
     *
     * @param key 键
     * @param value 值
     * @param <T> 值的类型
     */
    public <T> void put(Key<T> key, T value) {
        checkValue(key, value);
        List<Object> values = store.computeIfAbsent(key.name(), k -> new ArrayList<>(1));
        values.add(value);
    }
    
    /**
     * 移除键对应的所有值
     *
     * @param key 键
     * @param <T> 值的类型
     */
    public <T> void remove(Key<T> key) {
        store.remove(key.name());
    }
    
    /**
     * 获取所有键的集合
     *
     * @return 键的集合
     */
    public Set<String> keys() {
        return Collections.unmodifiableSet(store.keySet());
    }
    
    /**
     * 检查是否包含指定的键
     *
     * @param key 键
     * @param <T> 值的类型
     * @return 如果包含该键返回true，否则返回false
     */
    public <T> boolean containsKey(Key<T> key) {
        return store.containsKey(key.name());
    }
    
    /**
     * 移除键对应的所有值并返回第一个值
     *
     * @param key 键
     * @param <T> 值的类型
     * @return 被移除的第一个值，如果键不存在返回null
     */
    @SuppressWarnings("unchecked")
    public <T> T removeAll(Key<T> key) {
        List<Object> values = store.remove(key.name());
        if (values == null || values.isEmpty()) {
            return null;
        }
        return (T) values.get(0);
    }
    
    /**
     * 丢弃键对应的所有值
     *
     * @param key 键
     * @param <T> 值的类型
     */
    public <T> void discardAll(Key<T> key) {
        store.remove(key.name());
    }
    
    /**
     * 合并另一个元数据实例
     *
     * @param other 另一个元数据实例
     */
    public void merge(Metadata other) {
        if (other == null) {
            return;
        }
        for (Map.Entry<String, List<Object>> entry : other.store.entrySet()) {
            String key = entry.getKey();
            List<Object> otherValues = entry.getValue();
            if (otherValues != null && !otherValues.isEmpty()) {
                List<Object> values = store.computeIfAbsent(key, k -> new ArrayList<>());
                values.addAll(otherValues);
            }
        }
    }
    
    /**
     * 检查值是否与键的类型匹配
     *
     * @param key 键
     * @param value 值
     * @param <T> 值的类型
     * @throws IllegalArgumentException 如果值与键的类型不匹配
     */
    private <T> void checkValue(Key<T> key, T value) {
        if (value == null) {
            throw new NullPointerException("value cannot be null");
        }
        if (key.isBinary() && !(value instanceof byte[])) {
            throw new IllegalArgumentException("value must be byte[] for binary key");
        }
        if (!key.isBinary() && !(value instanceof String)) {
            throw new IllegalArgumentException("value must be String for ASCII key");
        }
        if (!key.isBinary() && value instanceof String) {
            String strValue = (String) value;
            if (!StandardCharsets.US_ASCII.newEncoder().canEncode(strValue)) {
                throw new IllegalArgumentException("value must be ASCII string");
            }
        }
    }
} 